package com.sunzm.flink.datastream.java.state;

import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

/**
 * JavaStateDemo
 *
 * @author Administrator
 * @version 1.0
 * @date 2021-06-22 20:48
 */
public class JavaStateDemo {
    private static boolean isLocal = true;

    public static void main(String[] args) throws Exception {
        //获取执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        if (isLocal) {
            env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(new Configuration());
        }

        //从数据源获取数据
        DataStream<String> dataStream = env.socketTextStream("82.156.210.70", 9999);

        KeyedStream<Tuple2<String, Double>, String> keyedDS = dataStream.flatMap(new FlatMapFunction<String, Tuple2<String, Double>>() {

            @Override
            public void flatMap(String value, Collector<Tuple2<String, Double>> out) throws Exception {

                if (StringUtils.isNotBlank(value)) {
                    //截取字符串
                    String[] fileds = value.split(",");
                    String name = StringUtils.trim(fileds[0]);
                    Double score = Double.parseDouble(StringUtils.trim(fileds[1]));

                    out.collect(Tuple2.of(name, score));
                }

            }
        }).keyBy(value -> value.f0);

        DataStream<Tuple2<String, Double>> resultDS = keyedDS.flatMap(new RichFlatMapFunction<Tuple2<String, Double>, Tuple2<String, Double>>() {

            /**
             * 记录学生的人数和总成绩
             */
            private ValueState<Tuple2<Integer, Double>> countAndSumState = null;

            /**
             * 初始化状态
             * 生命周期方法，只在启动时调用一次
             * @param parameters
             * @throws Exception
             */
            @Override
            public void open(Configuration parameters) throws Exception {
                countAndSumState = getRuntimeContext().getState(
                        new ValueStateDescriptor<Tuple2<Integer, Double>>("countAndSumState",
                                TypeInformation.of(new TypeHint<Tuple2<Integer, Double>>() {
                                })
                        )
                );
            }

            /**
             * 每来一条数据，都会调用一次，我们的业务逻辑主要就写在这个方法中
             * @param tp
             * @param collector
             * @throws Exception
             */
            @Override
            public void flatMap(Tuple2<String, Double> tp, Collector<Tuple2<String, Double>> collector) throws Exception {
                //line格式： zs, 99.5

                //得到名字和成绩字段的值
                String name = tp.f0;
                Double score = tp.f1;

                //获取状态中记录的学生个数和总成绩
                Tuple2<Integer, Double> countAndSum = countAndSumState.value();

                //第一次获取的时候，状态中可能没有数据, 需要给一个初始值
                if (countAndSum == null) {
                    countAndSum = Tuple2.of(0, 0.0D);
                }

                //状态中的总人数
                Integer oldCount = countAndSum.f0;
                //状态中的总分数
                Double oldScore = countAndSum.f1;

                //计算新的总人数和总分数
                Integer newCount = oldCount + 1;
                Double newScore = oldScore + score;

                //更新状态数据
                countAndSumState.update(Tuple2.of(newCount, newScore));

                //计算平均值
                double avgScore = newScore / newCount;
                //输出结果
                collector.collect(Tuple2.of(name, avgScore));

            }

            /**
             * 清空状态
             * 生命周期方法，只在flink程序退出时执行一次
             * @throws Exception
             */
            @Override
            public void close() throws Exception {
                //程序退出不能清空状态，不然重启后状态数据就无法恢复了
            }
        });

        //打印输出结果
        resultDS.print();

        //启动程序
        String className = Thread.currentThread().getStackTrace()[1].getClassName();
        String simpleName = className.substring(className.lastIndexOf(".") + 1);
        env.execute(simpleName);
    }
}
